import pandas as pd
from peewee import *
from datetime import datetime
from typing import Dict, List, Tuple

class SQLTableCreator:
    def __init__(self, db_config):
        """
        初始化数据库连接

        参数:
            db_config (dict): 数据库配置字典，包含:
                - 'db_type': 数据库类型 ('mysql', 'postgresql', 'sqlite')
                - 'database': 数据库名称
                - 'host': 主机地址 (可选)
                - 'port': 端口号 (可选)
                - 'user': 用户名 (可选)
                - 'password': 密码 (可选)
        """
        self.db_config = db_config
        self.db = self._connect_db()

    def _connect_db(self):
        """根据配置创建数据库连接"""
        db_type = self.db_config.get('db_type', 'mysql').lower()

        if db_type == 'mysql':
            return MySQLDatabase(
                database=self.db_config['database'],
                host=self.db_config.get('host', '127.0.0.1'),
                port=self.db_config.get('port', 3306),
                user=self.db_config.get('user', 'root'),
                password=self.db_config.get('password', 'pwb15017402748')
            )
        elif db_type == 'postgresql':
            return PostgresqlDatabase(
                self.db_config['database'],
                host=self.db_config.get('host', 'localhost'),
                port=self.db_config.get('port', 5432),
                user=self.db_config.get('user', 'postgres'),
                password=self.db_config.get('password', '')
            )
        elif db_type == 'sqlite':
            return SqliteDatabase(self.db_config['database'])
        else:
            raise ValueError(f"不支持的数据库类型: {db_type}")

    def _infer_field_type(self, series):
        """
        根据Pandas Series推断Peewee字段类型

        参数:
            series: Pandas Series对象

        返回:
            Peewee字段类型
        """
        # 检查是否为空Series
        if series.empty:
            return TextField(null=True)

        # 获取第一个非空值
        sample = series.dropna()
        if len(sample) == 0:
            return TextField(null=True)

        sample_value = sample.iloc[0]


        # 检查是否为布尔类型
        if pd.api.types.is_bool_dtype(series.dtype):
            return BooleanField(null=True)
        # 检查是否为数值类型
        elif pd.api.types.is_numeric_dtype(series.dtype):
            return DecimalField(max_digits=20,decimal_places=7,null=True)

        # 检查是否为日期时间类型
        elif pd.api.types.is_datetime64_any_dtype(series.dtype):
            return DateTimeField(null=True)

        # 默认为文本类型
        else:
            max_length = series.astype(str).str.len().max()
            # 限制最大长度为255，或者根据实际数据设置更大值
            return TextField(null=True)

    def create_table_from_dataframe(self, df, table_name, primary_key=None, if_exists='fail'):
        """
        根据DataFrame创建数据库表

        参数:
            df: Pandas DataFrame对象
            table_name (str): 要创建的表名
            primary_key (str, optional): 主键列名
            if_exists (str): 表已存在时的处理方式 ('fail', 'replace', 'append')

        返回:
            Peewee Model类
        """
        if if_exists not in ['fail', 'replace', 'append']:
            raise ValueError("if_exists必须是'fail', 'replace'或'append'")

        # 动态创建Model类
        attrs = {}
        field_definitions = {}

        for col in df.columns:
            field_type = self._infer_field_type(df[col])
            field_definitions[col] = field_type
            attrs[col] = field_type

        # 设置主键
        if primary_key and primary_key in df.columns:
            attrs[primary_key] = field_definitions[primary_key]
            attrs['_meta'] = {'primary_key': primary_key}
        else:
            # 如果没有指定主键，自动添加一个自增ID字段
            attrs['id'] = AutoField(primary_key=True)

        # 创建Model类
        model_class = type(table_name, (Model,), attrs)
        model_class._meta.database = self.db

        # 处理表已存在的情况
        if self.db.table_exists(table_name):
            if if_exists == 'fail':
                raise ValueError(f"表 {table_name} 已存在")
            elif if_exists == 'replace':
                self.db.drop_tables([model_class])
                self.db.create_tables([model_class])
        else:
            self.db.create_tables([model_class])

        return model_class

    def close(self):
        """关闭数据库连接"""
        if self.db:
            self.db.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()


def create_table(db_config: Dict, df: pd.DataFrame,
                 table_name: str, primary_key: str = None,
                 index: List[Tuple] = None):
    print("开始建表过程")
    # 使用上下文管理器创建表
    with SQLTableCreator(db_config) as creator:
        # 创建表 (如果存在则替换)
        model = creator.create_table_from_dataframe(
            df=df,
            table_name=table_name,
            primary_key=primary_key,
            if_exists='replace'
        )

        # 验证表是否创建成功
        print("表结构:")
        for field_name, field in model._meta.fields.items():
            print(f"{field_name}: {field.__class__.__name__}")

        # 插入数据
        with creator.db.atomic():
            for _, row in df.iterrows():
                record = {}
                for col in df.columns:
                    record[col] = row[col] if not pd.isna(row[col]) else None
                model.create(**record)
    print("结束建表过程")
    print("建表成功")

